#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import torch
import numpy as np
import sys
import time
import os
from enum import Enum

# ROS2相关导入
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist

os.environ['YOLO_VERBOSE'] = 'False'

class TrackingState(Enum):
    """跟踪状态枚举"""
    STOPPED = 0      # 停止
    FORWARD = 1      # 前进  
    TURN_LEFT = 5    # 左旋 (原地左转)
    TURN_RIGHT = 6   # 右旋 (原地右转)

class GreenBallTracker(Node):
    def __init__(self, model_path, camera_device=0):
        """
        初始化绿色小球跟踪器
        """
        super().__init__('green_ball_tracker')
        
        self.model_path = model_path
        self.camera_device = camera_device
        
        # 初始化ROS2发布器
        print("🔄 正在初始化ROS2控制...")
        try:
            self.pub_vel = self.create_publisher(Twist, '/cmd_vel', 1)
            time.sleep(0.1)
            print("✅ ROS2控制初始化成功")
        except Exception as e:
            print(f"❌ ROS2控制初始化失败: {e}")
            raise
        
        # 加载YOLO模型
        print("🔄 正在加载绿色小球检测模型...")
        try:
            # 修改这里：指向yolov5的目录，而不是当前目录
            self.model = torch.hub.load('/root/yolov5-7.0', 'custom', 
                                      path=model_path, source='local', force_reload=True)
            self.model.conf = 0.5
            self.model.iou = 0.45
            print(f"✅ 模型加载成功: {model_path}")
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            raise
        
        # 初始化摄像头
        print(f"🔄 正在初始化摄像头...")
        self.cap = cv2.VideoCapture(camera_device)
        if not self.cap.isOpened():
            raise RuntimeError(f"❌ 无法打开摄像头: /dev/video{camera_device}")
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 854)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print(f"✅ 摄像头初始化成功，分辨率: {self.frame_width}x{self.frame_height}")
        
        # ==================== 控制参数 ====================
        self.forward_speed = 0.2       # 前进速度 (m/s)
        self.turn_speed = 0.5          # 旋转速度 (rad/s)
        # ================================================
        
        # 跟踪参数
        self.center_tolerance = 10.0  # 中心容忍度（百分比）
        self.min_ball_size = 5.0      # 最小球体尺寸（百分比）
        self.max_ball_size = 70.0     # 最大球体尺寸（百分比）
        
        # 状态机
        self.current_state = TrackingState.STOPPED
        self.state_start_time = time.time()
        self.detection_count = 0
        
        # 状态机定时器间隔（秒）
        self.state_update_interval = 0.1  # 10Hz
        self.last_state_update = time.time()
        
    def publish_motion(self, linear_x=0.0, angular_z=0.0):
        """发布运动指令"""
        twist = Twist()
        twist.linear.x = linear_x
        twist.angular.z = angular_z
        self.pub_vel.publish(twist)
        
    def detect_balls(self, frame):
        """
        检测当前帧中的绿色小球
        返回置信度最高的小球信息
        """
        results = self.model(frame)
        detections = results.pandas().xyxy[0]
        
        best_ball = None
        highest_confidence = 0.0
        
        for _, detection in detections.iterrows():
            confidence = detection['confidence']
            if confidence > 0.5 and confidence > highest_confidence:
                # 计算相对坐标
                xmin, ymin = detection['xmin'], detection['ymin']
                xmax, ymax = detection['xmax'], detection['ymax']
                
                center_x = ((xmin + xmax) / 2 / self.frame_width) * 100
                center_y = ((ymin + ymax) / 2 / self.frame_height) * 100
                width = ((xmax - xmin) / self.frame_width) * 100
                height = ((ymax - ymin) / self.frame_height) * 100
                
                best_ball = {
                    'confidence': confidence,
                    'center_x': center_x,
                    'center_y': center_y,
                    'width': width,
                    'height': height
                }
                highest_confidence = confidence
        
        return best_ball
    
    def decide_next_state(self, ball_info):
        """
        根据小球位置决定下一个状态
        """
        if ball_info is None:
            return TrackingState.STOPPED
        
        center_x = ball_info['center_x']
        width = ball_info['width']
        height = ball_info['height']
        
        # 检查球体大小是否在合理范围内
        ball_size = (width + height) / 2
        if ball_size < self.min_ball_size or ball_size > self.max_ball_size:
            return TrackingState.STOPPED
        
        # 计算与画面中心的偏差
        center_offset = center_x - 50.0  # 50%是画面中心
        
        # 决策逻辑
        if abs(center_offset) <= self.center_tolerance:
            # 小球在中心区域
            if ball_size < 30.0:  # 球体较小，需要前进
                return TrackingState.FORWARD
            else:  # 球体较大，停止
                return TrackingState.STOPPED
        elif center_offset < -self.center_tolerance:
            # 小球在左侧，需要左旋
            return TrackingState.TURN_LEFT
        else:
            # 小球在右侧，需要右旋
            return TrackingState.TURN_RIGHT
    
    def execute_state(self, state):
        """
        执行状态对应的动作
        """
        if state == TrackingState.STOPPED:
            self.publish_motion(0.0, 0.0)  # 停止
        elif state == TrackingState.FORWARD:
            self.publish_motion(self.forward_speed, 0.0)  # 前进
        elif state == TrackingState.TURN_LEFT:
            self.publish_motion(0.0, self.turn_speed)  # 左旋（原地左转）
        elif state == TrackingState.TURN_RIGHT:
            self.publish_motion(0.0, -self.turn_speed)  # 右旋（原地右转）
    
    def transition_to(self, new_state):
        """状态转换"""
        if new_state != self.current_state:
            print(f"状态转换: {self.current_state.name} -> {new_state.name}")
            self.current_state = new_state
            self.state_start_time = time.time()
            self.execute_state(new_state)
    
    def state_machine_update(self, ball_info):
        """
        状态机更新
        """
        current_time = time.time()
        
        # 检查是否到了状态更新时间
        if current_time - self.last_state_update < self.state_update_interval:
            return
        
        self.last_state_update = current_time
        
        # 决定下一个状态
        next_state = self.decide_next_state(ball_info)
        
        # 状态转换
        self.transition_to(next_state)
    
    def run(self):
        """
        主运行循环
        """
        print("\n🚀 开始绿色小球跟踪...")
        print("📝 状态说明: STOPPED=停止, FORWARD=前进, TURN_LEFT=左旋, TURN_RIGHT=右旋")
        print("⌨️  按 Ctrl+C 退出\n")
        
        frame_count = 0
        last_fps_time = time.time()
        
        try:
            while rclpy.ok():
                # 读取摄像头帧
                ret, frame = self.cap.read()
                if not ret:
                    print("❌ 无法读取摄像头数据")
                    continue
                
                frame_count += 1
                
                # 检测绿色小球
                ball_info = self.detect_balls(frame)
                
                # 输出检测信息
                if ball_info:
                    self.detection_count += 1
                    print(f"[{self.detection_count:04d}] 🎯 置信度: {ball_info['confidence']:.1%} | "
                          f"中心: ({ball_info['center_x']:.1f}%, {ball_info['center_y']:.1f}%) | "
                          f"尺寸: {ball_info['width']:.1f}% x {ball_info['height']:.1f}% | "
                          f"状态: {self.current_state.name}")
                else:
                    # 无检测时也要更新状态（停止）
                    if frame_count % 30 == 0:  # 每30帧输出一次
                        print(f"[----] ❌ 未检测到小球 | 状态: {self.current_state.name}")
                
                # 状态机更新
                self.state_machine_update(ball_info)
                
                # FPS统计
                current_time = time.time()
                if current_time - last_fps_time >= 5.0:
                    fps = frame_count / (current_time - last_fps_time)
                    print(f"📊 处理速度: {fps:.1f} FPS | 累计检测: {self.detection_count} 次")
                    frame_count = 0
                    last_fps_time = current_time
                
                # 处理ROS2回调
                rclpy.spin_once(self, timeout_sec=0.001)
                
        except KeyboardInterrupt:
            print("\n🛑 用户停止跟踪")
        except Exception as e:
            print(f"\n💥 程序运行出错: {str(e)}")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """
        清理资源
        """
        print("🧹 正在清理资源...")
        
        # 停止小车
        for _ in range(3):
            self.publish_motion(0.0, 0.0)  # 停止
            time.sleep(0.05)
        
        # 释放摄像头
        if hasattr(self, 'cap') and self.cap:
            self.cap.release()
        
        cv2.destroyAllWindows()
        print(f"✅ 资源清理完成，总共检测到 {self.detection_count} 次绿色小球")

def main():
    model_path = "/root/yolov5-7.0/model/best.pt"
    
    print("🤖 绿色小球跟踪控制程序 - ROS2版本")
    print("=" * 50)
    print(f"📁 模型路径: {os.path.abspath(model_path)}")
    print(f"📁 工作目录: {os.getcwd()}")
    print("🎮 控制参数:")
    print("  - 前进速度: 0.2 m/s")
    print("  - 旋转速度: 0.2 rad/s")
    print("  - 中心容忍度: ±10%")
    print("  - 状态更新频率: 10Hz")
    print("📋 控制方式说明:")
    print("  - STOPPED: 停止 (linear=0, angular=0)")
    print("  - FORWARD: 前进 (linear=0.2, angular=0)") 
    print("  - TURN_LEFT: 左旋 (linear=0, angular=0.2)")
    print("  - TURN_RIGHT: 右旋 (linear=0, angular=-0.2)")
    
    # 初始化ROS2
    rclpy.init()
    
    try:
        tracker = GreenBallTracker(model_path, camera_device=0)
        tracker.run()
        
    except FileNotFoundError as e:
        print(f"📁 {str(e)}")
        print("💡 请确保模型文件 best.pt 在 model/ 目录中")
        return 1
    except Exception as e:
        print(f"💥 程序运行出错: {str(e)}")
        return 1
    finally:
        if rclpy.ok():
            rclpy.shutdown()
    
    return 0

if __name__ == "__main__":
    exit(main())